from typing import Tuple
import gym
from gym import spaces
from typing import Any

# import Any


class SingleMaze(gym.Env):
    def __init__(self, maze_size=6):
        self.maze_size = maze_size
        # Discrete(2) # {0, 1}
        self.action_space = spaces.Discrete(2)
        self.state_position = 0
        return

    def step(self, action: Any):
        assert self.action_space.contains(action), "%r (%s) invalid" % (
            action,
            type(action),
        )
        # 状态转移
        reward = 0
        done = False
        if self.state_position[0] == self.maze_size - 1:
            done = True
            self.state_position[0] = self.maze_size - 1
            reward = 10
        else:
            if action == 0:
                # 左移动，如果不在第一个位置，则向左移动，否则不动
                reward = -0.5
                if self.state_position[0] != 0:
                    self.state_position[0] -= 1
            else:
                self.state_position[0] += 1
        return self.state_position, reward, done

    def reset(self) -> Tuple[Any, dict]:
        # 把初始位置设置为中间位置
        self.state_position = [0, 0]
        return [0, 0]
